from EnvWrapper.BaseWrapper import BaseWrapper
from pathlib import Path
import sys
base_dir = Path(__file__).resolve().parent.parent.parent
sys.path.append(str(base_dir))
from env.chooseenv import make
env = make("cliffwalking")

import pygame
pygame.init()
screen = pygame.display.set_mode(env.grid.size)
clock = pygame.time.Clock()


class cliffwalking(BaseWrapper):
    def __init__(self):
        self.env = env
        super().__init__(self.env)

    def get_actionspace(self):
        return self.env.action_dim

    def get_observationspace(self):
        return int(self.env.input_dimension)

    def step(self, action, train=True):
        '''
        return: next_state, reward, done, _, _
        '''

        next_state, reward, done, _, info = self.env.step(action)
        reward = reward[0]
        return next_state, reward, done, _, info

    def reset(self):
        state = self.env.reset()
        return state

    def close(self):
        # pygame.quit()
        pass

    def set_seed(self, seed=None):
        pass

    def make_render(self):
        pygame.surfarray.blit_array(screen, self.env.render_board().transpose(1, 0, 2))
        pygame.display.flip()
        # fps=200
        # self.clock.tick(fps)

